# ================== Regan 攻击层 =====================
import os
import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image
from .RegenMark.simulator import test_Regen_simulator  # 假设推理类在 test_Regen_simulator.py



class Regen(nn.Module):
    def __init__(self, model_path="RegenMark/simulator/best_model.pth", device=None):
        super(Regen, self).__init__()
        self.device = device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
        # 初始化推理器
        self.inferencer = test_Regen_simulator.TwoStageInference(model_path, self.device)
        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
        ])

    def forward(self, image_cover_mask):
        image = image_cover_mask[0]
        noised_image = torch.zeros_like(image)
        for i in range(image.shape[0]):
            # 转为PIL
            single_image = ((image[i].clamp(-1, 1).permute(1, 2, 0) + 1) / 2 * 255).add(0.5).clamp(0, 255).to('cpu', torch.uint8).numpy()
            pil_img = Image.fromarray(single_image)
            # 推理去水印
            result = self.inferencer.inference_single(pil_img)
            restored = result['restored_image']
            # 转回tensor
            noised_image[i] = self.transform(restored).to(image.device)
        return noised_image